package sshex

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"sync"
	"time"

	"gitee.com/shanyou/sshplus/rlog"
	"github.com/oklog/run"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
)

// Option is a functional option handler for Server.
type Option func(*SSHServer) error

// SSHServer ssh 服务实现
type SSHServer struct {
	mu      sync.RWMutex
	canncel context.CancelFunc
	ctx     context.Context
	conns   sync.Map
	connWg  sync.WaitGroup

	Version                    string
	Logger                     logrus.FieldLogger
	HostSigners                []ssh.Signer
	IdleTimeout                time.Duration // connection timeout when no activity, none if empty
	MaxTimeout                 time.Duration // absolute connection timeout, none if empty
	Listener                   net.Listener
	ChannelHandlers            map[string]ChannelHandler
	RequestHandlers            map[string]RequestHandler
	ConnectionFailedCallback   ConnectionFailedCallback
	PasswordHandler            PasswordHandler // password authentication handler
	KeyboardInteractiveHandler func(conn ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error)
	PublicKeyHandler           func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error)
}

func defaultOption(srv *SSHServer) error {
	return nil
}

func NewSSHServer(version string, options ...Option) *SSHServer {
	l := rlog.Logger.WithField("com", "ssh-server")
	srv := &SSHServer{
		Logger:  l,
		Version: version,
	}
	defaultOption(srv)
	for _, opt := range options {
		srv.SetOption(opt)
	}
	return srv
}

// SetOption runs a functional option against the server.
func (srv *SSHServer) SetOption(option Option) error {
	// NOTE: there is a potential race here for any option that doesn't call an
	// internal method. We can't actually lock here because if something calls
	// (as an example) AddHostKey, it will deadlock.

	//srv.mu.Lock()
	//defer srv.mu.Unlock()

	return option(srv)
}

func (srv *SSHServer) config(ctx *ConnContext) *ssh.ServerConfig {
	srv.mu.RLock()
	defer srv.mu.RUnlock()

	config := &ssh.ServerConfig{}
	for _, signer := range srv.HostSigners {
		config.AddHostKey(signer)
	}
	if srv.Version != "" {
		config.ServerVersion = "SSH-2.0-" + srv.Version
	}

	if srv.PasswordHandler != nil {
		config.PasswordCallback = func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			ctx.ApplyConnMetadata(conn)
			if ok := srv.PasswordHandler(ctx, string(password)); !ok {
				return ctx.Permissions(), fmt.Errorf("permission denied")
			}
			return ctx.Permissions(), nil
		}
	}

	if srv.KeyboardInteractiveHandler != nil {
		config.KeyboardInteractiveCallback = func(conn ssh.ConnMetadata, challenger ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
			ctx.ApplyConnMetadata(conn)
			return srv.KeyboardInteractiveHandler(conn, challenger)
		}
	}

	if srv.PublicKeyHandler != nil {
		config.PublicKeyCallback = func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			ctx.ApplyConnMetadata(conn)
			perm, err := srv.PublicKeyHandler(conn, key)
			if err == nil {
				ctx.SetValue(ContextKeyPublicKey, key)
			}

			return perm, err
		}
	}
	return config
}

func (srv *SSHServer) AddHostKeyCertWithBytes(keyBytes, certBytes []byte) error {
	block, _ := pem.Decode(keyBytes)
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	singer, err := ssh.NewSignerFromSigner(key)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	pubKey, _, _, _, err := ssh.ParseAuthorizedKey(certBytes)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	certPub, ok := pubKey.(*ssh.Certificate)
	if !ok {
		return fmt.Errorf("cert format error")
	}
	certSigner, err := ssh.NewCertSigner(certPub, singer)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	srv.AddHostKey(certSigner)
	return nil
}

func (srv *SSHServer) AddHostKeyWithBytes(keyBytes []byte) error {
	block, _ := pem.Decode(keyBytes)
	key, err := x509.ParseECPrivateKey(block.Bytes)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	singer, err := ssh.NewSignerFromSigner(key)
	if err != nil {
		return fmt.Errorf("add key to agent error %v", err)
	}
	srv.AddHostKey(singer)
	return nil
}

// AddHostKey adds a private key as a host key. If an existing host key exists
// with the same algorithm, it is overwritten. Each server config must have at
// least one host key.
func (srv *SSHServer) AddHostKey(key ssh.Signer) {
	srv.mu.Lock()
	defer srv.mu.Unlock()

	// these are later added via AddHostKey on ServerConfig, which performs the
	// check for one of every algorithm.

	// This check is based on the AddHostKey method from the x/crypto/ssh
	// library. This allows us to only keep one active key for each type on a
	// server at once. So, if you're dynamically updating keys at runtime, this
	// list will not keep growing.
	for i, k := range srv.HostSigners {
		if k.PublicKey().Type() == key.PublicKey().Type() {
			srv.HostSigners[i] = key
			return
		}
	}

	srv.HostSigners = append(srv.HostSigners, key)
}

func (srv *SSHServer) EnsureHostSigner() error {
	srv.mu.Lock()
	defer srv.mu.Unlock()

	if len(srv.HostSigners) == 0 {
		key, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			return err
		}
		singer, err := ssh.NewSignerFromKey(key)
		if err != nil {
			return err
		}
		srv.HostSigners = append(srv.HostSigners, singer)
	}
	return nil
}

func (srv *SSHServer) Serve(ctx context.Context, l net.Listener) error {
	srv.Logger.Debug("begin serve ssh server")
	srv.SetContext(ctx)
	srv.Listener = l
	defer l.Close()
	if err := srv.EnsureHostSigner(); err != nil {
		return err
	}
	for {
		conn, e := l.Accept()
		if e != nil {
			select {
			case <-srv.ctx.Done():
				return ErrServerClosed
			default:
				srv.Logger.WithError(e).Error("serve connect err")
				continue
			}
		}
		go srv.HandleConn(conn)
	}
}

func (srv *SSHServer) SetContext(ctx context.Context) {
	srv.ctx, srv.canncel = context.WithCancel(ctx)
}

// HandleConn must SetContext and EnsureHostSigner before run
func (srv *SSHServer) HandleConn(newConn net.Conn) {
	ctx, cancel := NewConnContext(srv.ctx, newConn)
	conn := &serverConn{
		Conn:          newConn,
		idleTimeout:   srv.IdleTimeout,
		closeCanceler: cancel,
	}
	if srv.MaxTimeout > 0 {
		conn.maxDeadline = time.Now().Add(srv.MaxTimeout)
	}
	//create ssh connection
	srv.Logger.Debugf("create ssh server conn for %v", conn)
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, srv.config(ctx))
	if err != nil {
		srv.Logger.Warningf("create ssh connection error : %s", err.Error())
		if srv.ConnectionFailedCallback != nil {
			srv.ConnectionFailedCallback(conn, err)
		}
		return
	}
	srv.trackConn(sshConn, true)
	defer func() {
		srv.Logger.Debugf("close ssh server conn for %v", conn)
		conn.Close()
		srv.trackConn(sshConn, false)
	}()
	ctx.ApplyConnMetadata(sshConn)
	var g run.Group
	{
		g.Add(func() error {
			for req := range reqs {
				srv.Logger.Debugf("receive request type: %s", req.Type)
				handler := srv.RequestHandlers[req.Type]
				if handler == nil {
					handler = srv.RequestHandlers["default"]
				}
				if handler == nil {
					req.Reply(false, nil)
					continue
				}
				/*reqCtx, cancel := context.WithCancel(ctx)
				defer cancel() */
				reply, ret, payload := handler(ctx, sshConn, req)
				if reply {
					req.Reply(ret, payload)
				}
			}
			return nil
		}, func(err error) {})

		g.Add(func() error {
			for ch := range chans {
				handler := srv.ChannelHandlers[ch.ChannelType()]
				if handler == nil {
					handler = srv.ChannelHandlers["default"]
				}
				if handler == nil {
					ch.Reject(ssh.UnknownChannelType, "unsupported channel type")
					continue
				}
				go handler(ctx, sshConn, ch)
			}
			return nil
		}, func(err error) {})
	}
	err = g.Run()
	if err != nil {
		srv.Logger.Debugf("handle conn err: %s", err.Error())
	}
}

func (srv *SSHServer) Shutdown() error {
	srv.mu.Lock()
	unset := srv.Listener == nil
	srv.mu.Unlock()
	if unset {
		return errors.New("not started yet")
	}

	srv.mu.Lock()
	l := srv.Listener
	cancel := srv.canncel
	srv.mu.Unlock()
	cancel()
	return l.Close()
}

func (srv *SSHServer) Wait() error {
	srv.mu.Lock()
	unset := srv.Listener == nil
	srv.mu.Unlock()
	if unset {
		return errors.New("not started yet")
	}

	finished := make(chan struct{}, 1)
	go func() {
		srv.connWg.Wait()
		finished <- struct{}{}
	}()

	select {
	case <-srv.ctx.Done():
		return srv.ctx.Err()
	case <-finished:
		return nil
	}
}

// trackConn 存储connection链接
func (srv *SSHServer) trackConn(c *ssh.ServerConn, add bool) {
	if add {
		srv.conns.LoadOrStore(c, struct{}{})
		srv.connWg.Add(1)
	} else {
		srv.conns.Delete(c)
		srv.connWg.Done()
	}
}
